package storemanmpc

import (
	"errors"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	"github.com/wanchain/go-mpc/storeman/schnorrmpc"
	"github.com/wanchain/go-mpc/storeman/schnorrmpcbn"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"github.com/wanchain/go-mpc/storeman/storemanmpc/step"
)

func reqSignMpcSchnorr(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCurCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	reqMpc := step.CreateRequestMpcStepUni(int(totalNum), peerCurCount, mpcprotocol.MpcSignLeaderUni)
	reqMpc.SetWaiting(int(peerCurCount))

	mpcReady := step.CreateMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(0)

	return generateTxSignMpcSchnorr(mpc, reqMpc, mpcReady, peerCurCount, curveType)
}

func ackSignMpcSchnorr(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCurCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	ackMpc := step.CreateAckMpcStepUni(int(totalNum), mpcprotocol.MpcSignPeerUni)
	ackMpc.SetWaiting(0)

	mpcReady := step.CreateGetMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(1)

	return generateTxSignMpcSchnorr(mpc, ackMpc, mpcReady, peerCurCount, curveType)
}

func generateTxSignMpcSchnorr(mpc *MpcContext, firstStep MpcStepFunc, readyStep MpcStepFunc, peerCurCount uint16, curveType uint8) (*MpcContext, error) {
	log.SyslogInfo("generateTxSignMpcSchnorr begin", "ctx", mpc.ContextID)

	sigNum := len(mpc.GetInput().HashData)
	result := mpc.mpcResult
	_, grpIdString, err := osmconf.GetGrpId(result)
	if err != nil {
		log.SyslogErr("generateTxSignMpcSchnorr", "GetGrpId error", err.Error(), "ctx", mpc.ContextID)
		return nil, err
	}

	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	threshold, _ := osmconf.GetOsmConf().GetThresholdNum(grpIdString)
	if threshold < uint16(1) {
		log.SyslogErr("generateTxSignMpcSchnorr", "threshold error", threshold, "ctx", mpc.ContextID)
		return nil, errors.New("invalid threshold")
	}
	degree := threshold - 1

	// 2
	approveNumbStep := step.CreateMpcApprovedNumStepUni(int(totalNum))
	approveNumbStep.SetWaiting(1)

	// 3
	cmStep := step.CreateMpcPolycmStepUni(int(totalNum), sigNum)
	cmStep.SetWaiting(int(peerCurCount + 1))

	// 4
	skShare := step.CreateMpcRSKShareStepUni(int(degree), int(totalNum), sigNum)
	skShare.SetWaiting(int(peerCurCount + 1)) // not broadcast, only need receive peerCurCount data.

	// 5
	skJudgeStep := step.CreateMpcRSkJudgeStepUni(int(totalNum), sigNum)
	// only handle the first Rsk challenge or (timeout no challenge)
	skJudgeStep.SetWaiting(1)
	//skJudgeStep.SetWaiting(int(peerCurCount + 1))

	// 6
	// add rrcvInter step
	rrcvInterStep := step.CreateMpcRRcvInterStepUni(int(totalNum))
	rrcvInterStep.SetWaiting(int(peerCurCount + 1))

	// 7
	// add rrcvInter judge step
	rrcvJudgeStep := step.CreateMpcRRcvJudgeStepUni(int(totalNum), sigNum)
	rrcvJudgeStep.SetWaiting(int(peerCurCount + 1))

	// 8
	RCommitStep := step.CreateMpcRCommitStepUni(int(totalNum), sigNum)
	RCommitStep.SetWaiting(int(peerCurCount + 1))

	// 9
	RStep := step.CreateMpcRStepUni(int(totalNum), sigNum)
	//since RPKShare is used to check the content of sshare, so it should collect all node's RPK.
	//otherwise, when check sshare content system can not find the RPKShare.
	rrcvJudgeStep.SetWaiting(int(peerCurCount + 1))
	//RStep.SetWaiting(int(threshold))

	// 10
	resultKeys := mpcprotocol.GetPreSetKeyArrUni(mpcprotocol.MpcS, sigNum)
	SStep := step.CreateMpcSStepUni(int(totalNum), sigNum, resultKeys)
	//SStep.SetWaiting(int(threshold))
	SStep.SetWaiting(int(peerCurCount + 1))

	// 11
	sshareJudgeStep := step.CreateMpcSSahreJudgeStepUni(int(totalNum))
	// only handle the first sshare challenge or (timeout no challenge)
	//sshareJudgeStep.SetWaiting(1)
	//sshareJudgeStep.SetWaiting(int(peerCurCount + 1))
	sshareJudgeStep.SetWaiting(int(peerCurCount))

	// 12
	ackRSStep := step.CreateAckMpcRSStepUni(int(totalNum), sigNum)
	//ackRSStep.SetWaiting(int(peerCurCount + 1))
	ackRSStep.SetWaiting(int(peerCurCount))

	mpc.setMpcStep(firstStep,
		readyStep,
		approveNumbStep,
		cmStep,
		skShare,
		skJudgeStep,
		rrcvInterStep,
		rrcvJudgeStep,
		RCommitStep,
		RStep,
		SStep,
		sshareJudgeStep,
		ackRSStep)

	for stepId, stepItem := range mpc.MpcSteps {
		//stepItem.SetWaiting(len(mpc.peers) + 1)
		stepItem.SetWaitAll(false)
		stepItem.SetStepId(stepId)

		stepItem.SetMpcContextId(mpc.ContextID)

		stepItem.SetMpcContext(mpc)
	}

	var schnorrMpcer mpcprotocol.SchnorrMPCer
	switch int(curveType) {
	case mpcprotocol.SK256Curve:
		log.SyslogInfo("generateTxSignMpcSchnorr Using SK256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	case mpcprotocol.BN256Curve:
		log.SyslogInfo("generateTxSignMpcSchnorr Using BN256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpcbn.NewBnSchnorrMpc()
	default:
		log.SyslogInfo("generateTxSignMpcSchnorr Using SK256 curve (default...)", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	}

	mpc.SetSchnorrMPCer(schnorrMpcer)
	mpc.SetStepSchnorrMPCer()

	mpc.SetStepCurveType(uint16(curveType))

	for _, step := range mpc.MpcSteps {
		msgGens := step.GetMsgGens()
		if msgGens != nil {
			for _, msgGen := range msgGens {
				msgGen.SetSchnorrMpcer(schnorrMpcer)
				msgGen.SetMpcContext(mpc)
			}
		}
	}
	return mpc, nil
}

func reqSignMpcSchnorr340(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCurCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	reqMpc := step.CreateRequestMpcStepUni(int(totalNum), peerCurCount, mpcprotocol.MpcSignLeaderUni)
	reqMpc.SetWaiting(int(peerCurCount))

	mpcReady := step.CreateMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(0)

	return generateTxSignMpcSchnorr340(mpc, reqMpc, mpcReady, peerCurCount, curveType)
}

func ackSignMpcSchnorr340(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCurCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	ackMpc := step.CreateAckMpcStepUni(int(totalNum), mpcprotocol.MpcSignPeerUni)
	ackMpc.SetWaiting(0)

	mpcReady := step.CreateGetMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(1)

	return generateTxSignMpcSchnorr340(mpc, ackMpc, mpcReady, peerCurCount, curveType)
}

func generateTxSignMpcSchnorr340(mpc *MpcContext, firstStep MpcStepFunc, readyStep MpcStepFunc, peerCurCount uint16, curveType uint8) (*MpcContext, error) {
	log.SyslogInfo("generateTxSignMpcSchnorr340 begin", "ctx", mpc.ContextID)

	sigNum := len(mpc.GetInput().HashData)
	result := mpc.mpcResult
	_, grpIdString, err := osmconf.GetGrpId(result)
	if err != nil {
		log.SyslogErr("generateTxSignMpc", "GetGrpId error", err.Error(), "ctx", mpc.ContextID)
		return nil, err
	}

	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	threshold, _ := osmconf.GetOsmConf().GetThresholdNum(grpIdString)
	if threshold < uint16(1) {
		log.SyslogErr("generateTxSignMpc", "threshold error", threshold, "ctx", mpc.ContextID)
		return nil, errors.New("invalid threshold")
	}
	degree := threshold - 1

	//
	approveNumbStep := step.CreateMpcApprovedNumStepUni(int(totalNum))
	approveNumbStep.SetWaiting(1)

	// 2
	cmStep := step.CreateMpcPolycmStepUni(int(totalNum), sigNum)
	cmStep.SetWaiting(int(peerCurCount + 1))

	// 3
	skShare := step.CreateMpcRSKShareStepUni(int(degree), int(totalNum), sigNum)
	skShare.SetWaiting(int(peerCurCount + 1)) // not broadcast, only need receive peerCurCount data.

	// 4
	skJudgeStep := step.CreateMpcRSkJudgeStepUni(int(totalNum), sigNum)
	// only handle the first Rsk challenge or (timeout no challenge)
	skJudgeStep.SetWaiting(1)
	//skJudgeStep.SetWaiting(int(peerCurCount + 1))

	// 5
	// add rrcvInter step
	rrcvInterStep := step.CreateMpcRRcvInterStepUni(int(totalNum))
	rrcvInterStep.SetWaiting(int(peerCurCount + 1))

	// 6
	// add rrcvInter judge step
	rrcvJudgeStep := step.CreateMpcRRcvJudgeStepUni(int(totalNum), sigNum)
	rrcvJudgeStep.SetWaiting(int(peerCurCount + 1))

	RCommitStep := step.CreateMpcRCommitStepUni(int(totalNum), sigNum)
	RCommitStep.SetWaiting(int(peerCurCount + 1))

	// 7
	RStep := step.CreateMpcRStepUni(int(totalNum), sigNum)
	//since RPKShare is used to check the content of sshare, so it should collect all node's RPK.
	//otherwise, when check sshare content system can not find the RPKShare.
	rrcvJudgeStep.SetWaiting(int(peerCurCount + 1))
	//RStep.SetWaiting(int(threshold))

	// 8
	resultKeys := mpcprotocol.GetPreSetKeyArrUni(mpcprotocol.MpcS, sigNum)
	SStep := step.CreateMpcSStepUni(int(totalNum), sigNum, resultKeys)
	//SStep.SetWaiting(int(threshold))
	SStep.SetWaiting(int(peerCurCount + 1))

	// 9
	sshareJudgeStep := step.CreateMpcSSahreJudgeStepUni(int(totalNum))
	// only handle the first sshare challenge or (timeout no challenge)
	//sshareJudgeStep.SetWaiting(1)
	//sshareJudgeStep.SetWaiting(int(peerCurCount + 1))
	sshareJudgeStep.SetWaiting(int(peerCurCount))

	// 10
	ackRSStep := step.CreateAckMpcRSStepUni340(int(totalNum), sigNum)
	//ackRSStep.SetWaiting(int(peerCurCount + 1))
	ackRSStep.SetWaiting(int(peerCurCount))

	mpc.setMpcStep(firstStep,
		readyStep,
		approveNumbStep,
		cmStep,
		skShare,
		skJudgeStep,
		rrcvInterStep,
		rrcvJudgeStep,
		RCommitStep,
		RStep,
		SStep,
		sshareJudgeStep,
		ackRSStep)

	for stepId, stepItem := range mpc.MpcSteps {
		//stepItem.SetWaiting(len(mpc.peers) + 1)
		stepItem.SetWaitAll(false)
		stepItem.SetStepId(stepId)

		stepItem.SetMpcContextId(mpc.ContextID)

		stepItem.SetMpcContext(mpc)
	}

	var schnorrMpcer mpcprotocol.SchnorrMPCer
	switch int(curveType) {
	case mpcprotocol.SK256Curve:
		log.SyslogInfo("Using SK256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	case mpcprotocol.BN256Curve:
		log.SyslogInfo("Using BN256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpcbn.NewBnSchnorrMpc()
	default:
		log.SyslogInfo("Using SK256 curve (default...)", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	}

	mpc.SetSchnorrMPCer(schnorrMpcer)
	mpc.SetStepSchnorrMPCer()

	mpc.SetStepCurveType(uint16(curveType))

	for _, step := range mpc.MpcSteps {
		msgGens := step.GetMsgGens()
		if msgGens != nil {
			for _, msgGen := range msgGens {
				msgGen.SetSchnorrMpcer(schnorrMpcer)
				msgGen.SetMpcContext(mpc)
			}
		}
	}
	return mpc, nil
}

func reqSignMpcEcdsa(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCurCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	reqMpc := step.CreateRequestMpcStepUni(int(totalNum), peerCurCount, mpcprotocol.MpcSignLeaderUni)
	reqMpc.SetWaiting(int(peerCurCount))

	mpcReady := step.CreateMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(0)

	return generateTxSignMpcEcdsa(mpc, reqMpc, mpcReady, peerCurCount, curveType)
}

func ackSignMpcEcdsa(mpcID uint64, peers []mpcprotocol.PeerInfo, peerCount uint16, curveType uint8, alg uint8, input *mpcprotocol.SignInput, preSetValue ...MpcValue) (*MpcContext, error) {
	result := createMpcBaseMpcResult()
	result.InitializeValue(preSetValue...)
	result.SetMpcContextId(mpcID)

	mpc := createMpcContext(mpcID, peers, result, curveType)
	if mpc != nil {
		mpc.SetAlg(alg)
		mpc.SetInput(input)
	}

	_, grpIdString, _ := osmconf.GetGrpId(result)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpIdString)

	ackMpc := step.CreateAckMpcStepUni(int(totalNum), mpcprotocol.MpcSignPeerUni)
	ackMpc.SetWaiting(0)

	mpcReady := step.CreateGetMpcReadyStepUni(int(totalNum))
	mpcReady.SetWaiting(1)
	// wait 1
	return generateTxSignMpcEcdsa(mpc, ackMpc, mpcReady, peerCount, curveType)
}

func generateTxSignMpcEcdsa(mpc *MpcContext, firstStep MpcStepFunc, readyStep MpcStepFunc, peerCount uint16, curveType uint8) (*MpcContext, error) {
	log.SyslogInfo("generateTxSignMpcEcdsa begin", "ctx", mpc.ContextID)

	signNum := len(mpc.GetInput().HashData)

	// 2
	// get degree from config.
	_, grpId, err := osmconf.GetGrpId(mpc.mpcResult)
	if err != nil {
		return nil, err
	}
	degree, _ := osmconf.GetOsmConf().GetThresholdNumByCurve(grpId, curveType)
	totalNum, _ := osmconf.GetOsmConf().GetTotalNum(grpId)

	//
	approveNumbStep := step.CreateMpcApprovedNumStepUni(int(totalNum))
	approveNumbStep.SetWaiting(1)

	//
	log.SyslogInfo("generateTxSignMpcEcdsa", "degree", degree, "ctx", mpc.ContextID)
	JRJZ := step.CreateTXSignJR_JZ_StepUni(int(degree), int(totalNum), signNum)
	JRJZ.SetWaiting(int(peerCount + 1))

	// 3
	pointStepPreValueKeys := mpcprotocol.GetPreSetKeyArr(mpcprotocol.MpcSignA0, signNum)
	pointStepResultKeys := mpcprotocol.GetPreSetKeyArr(mpcprotocol.MpcSignAPoint, signNum)
	AGPoint := step.CreateMpcPoint_StepUni(int(totalNum), pointStepPreValueKeys, pointStepResultKeys)
	AGPoint.SetWaiting(int(peerCount + 1))

	// 4
	peerIntersaction := step.CreateMpcPeerIntersactionStepUni(int(totalNum))
	log.SyslogInfo("peerIntersaction", "peerCount", peerCount, "ctx", mpc.ContextID)
	peerIntersaction.SetWaiting(int(peerCount + 1))

	// 5
	peerGroupStep := step.CreateMpcPeerGroupStepUni(int(totalNum), signNum)
	log.SyslogInfo("peerGroupStep", "peerCount", peerCount, "ctx", mpc.ContextID)
	peerGroupStep.SetWaiting(int(peerCount + 1))

	// 6
	lagStepPreValueKeys := mpcprotocol.GetPreSetKeyArr(mpcprotocol.MpcSignARSeed, signNum)
	lagStepResultKeys := mpcprotocol.GetPreSetKeyArr(mpcprotocol.MpcSignARResult, signNum)
	ARLag := step.CreateTXSign_Lagrange_StepUni(int(totalNum), lagStepPreValueKeys, lagStepResultKeys)
	ARLag.SetWaiting(int(peerCount + 1))

	// 7
	TXSignLag := step.CreateTxSign_CalSignStepUni(int(totalNum), mpcprotocol.MpcTxSignResult, signNum)
	TXSignLag.SetWaiting(int(peerCount + 1))

	mpc.setMpcStep(firstStep,
		readyStep,
		approveNumbStep,
		JRJZ,
		AGPoint,
		peerIntersaction,
		peerGroupStep,
		ARLag,
		TXSignLag)

	for index, step := range mpc.MpcSteps {
		step.SetStepId(index)
		step.SetWaitAll(false)
		step.SetMpcContext(mpc)
		step.SetMpcContextId(mpc.ContextID)

		step.SetMpcContext(mpc)
	}

	var schnorrMpcer mpcprotocol.SchnorrMPCer
	switch int(curveType) {
	case mpcprotocol.SK256Curve:
		log.SyslogInfo("Using SK256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	case mpcprotocol.BN256Curve:
		log.SyslogInfo("Using BN256 curve", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpcbn.NewBnSchnorrMpc()
	default:
		log.SyslogInfo("Using SK256 curve (default...)", "ctx", mpc.ContextID)
		schnorrMpcer = schnorrmpc.NewSkSchnorrMpc()
	}

	mpc.SetSchnorrMPCer(schnorrMpcer)
	mpc.SetStepSchnorrMPCer()
	mpc.SetStepCurveType(uint16(curveType))

	for _, step := range mpc.MpcSteps {
		msgGens := step.GetMsgGens()
		if msgGens != nil {
			for _, msgGen := range msgGens {
				msgGen.SetSchnorrMpcer(schnorrMpcer)
				msgGen.SetContextId(mpc.ContextID)
				msgGen.SetMpcContext(mpc)
			}
		}
	}

	return mpc, nil
}
